#include <stdlib.h>
#include "path_planning/avl.h"

inline static int MIN(int a, int b) {
   return  (a < b) ? a : b;
}

inline static int MAX(int a, int b) {
   return  (a > b) ? a : b;
}

enum balance_t { LEFT_HEAVY = -1, BALANCED = 0, RIGHT_HEAVY = 1 };

enum height_effect_t { HEIGHT_NOCHANGE = 0, HEIGHT_CHANGE = 1 };

// Return true if the tree is too heavy on the left side
inline static int
LEFT_IMBALANCE(short bal) {
   return (bal < LEFT_HEAVY);
}

// Return true if the tree is too heavy on the right side
inline static int RIGHT_IMBALANCE(short bal) {
   return (bal > RIGHT_HEAVY);
}

AvlNode::AvlNode(Comparable* item) : myData(item), myBal(0)
{
   Reset();
}

AvlNode::~AvlNode(void) {
   if (mySubtree[LEFT])  delete  mySubtree[LEFT];
   if (mySubtree[RIGHT]) delete  mySubtree[RIGHT];
}

int AvlNode::RotateOnce(AvlNode* &root, dir_t dir)
{
   dir_t  otherDir = Opposite(dir);
   AvlNode* oldRoot = root;

   int  heightChange = (root->mySubtree[otherDir]->myBal == 0) ? HEIGHT_NOCHANGE : HEIGHT_CHANGE;

   root = oldRoot->mySubtree[otherDir];

   oldRoot->mySubtree[otherDir] = root->mySubtree[dir];
   root->mySubtree[dir] = oldRoot;

   oldRoot->myBal = -((dir == LEFT) ? --(root->myBal) : ++(root->myBal));

   return  heightChange;
}

int AvlNode::RotateTwice(AvlNode* &root, dir_t dir)
{
   dir_t  otherDir = Opposite(dir);
   AvlNode* oldRoot = root;
   AvlNode* oldOtherDirSubtree = root->mySubtree[otherDir];

   root = oldRoot->mySubtree[otherDir]->mySubtree[dir];

   oldRoot->mySubtree[otherDir] = root->mySubtree[dir];
   root->mySubtree[dir] = oldRoot;

   oldOtherDirSubtree->mySubtree[dir] = root->mySubtree[otherDir];
   root->mySubtree[otherDir] = oldOtherDirSubtree;

   // update balances equations
   root->mySubtree[LEFT]->myBal  = -MAX(root->myBal, 0);
   root->mySubtree[RIGHT]->myBal = -MIN(root->myBal, 0);
   root->myBal = 0;

   // A double rotation always shortens the overall height of the tree
   return  HEIGHT_CHANGE;
}

int AvlNode::ReBalance(AvlNode* &root) {
   int  heightChange = HEIGHT_NOCHANGE;

   if (LEFT_IMBALANCE(root->myBal)) {
      // Need a right rotation
      if (root->mySubtree[LEFT]->myBal  ==  RIGHT_HEAVY) {
         // RL rotation needed
         heightChange = RotateTwice(root, RIGHT);
      } else {
         // RR rotation needed
         heightChange = RotateOnce(root, RIGHT);
      }
   } else if (RIGHT_IMBALANCE(root->myBal)) {
      // Need a left rotation
      if (root->mySubtree[RIGHT]->myBal  ==  LEFT_HEAVY) {
         // LR rotation needed
         heightChange = RotateTwice(root, LEFT);
      } else {
         // LL rotation needed
         heightChange = RotateOnce(root, LEFT);
      }
   }

   return  heightChange;
}

cmp_t AvlNode::Compare(AvlKey key, cmp_t cmp) const
{
   switch (cmp) {
      case EQ_CMP :  // Standard comparison
         return  myData->Compare(key);

      case MIN_CMP :  // Find the minimal element in this tree
         return  (mySubtree[LEFT] == NULL) ? EQ_CMP : MIN_CMP;

      case MAX_CMP :  // Find the maximal element in this tree
         return  (mySubtree[RIGHT] == NULL) ? EQ_CMP : MAX_CMP;
   }
}

Comparable* AvlNode::Search(AvlKey key, AvlNode* root, cmp_t cmp)
{
   cmp_t result;
   while (root  &&  (result = root->Compare(key, cmp))) {
      root = root->mySubtree[(result < 0) ? LEFT : RIGHT];
   }
   return  (root) ? root->myData : NULL;
}

Comparable* AvlNode::Insert(Comparable* item, AvlNode* &root)
{
   int  change;
   return  Insert(item, root, change);
}

Comparable* AvlNode::PickNode(AvlNode* &root)
{
	AvlNode *aux=root;
   while(aux && aux->mySubtree[LEFT]){
		aux = aux->mySubtree[LEFT];
	}
   return Delete(aux->Key(), root, EQ_CMP);
}

Comparable* AvlNode::Delete(AvlKey key, AvlNode* &root, cmp_t cmp)
{
   int  change;
   return  Delete(key, root, change, cmp);
}

Comparable* AvlNode::Insert(Comparable*   item, AvlNode* &root, int &change)
{
   if (root == NULL) { 
      root = new AvlNode(item);
      change =  HEIGHT_CHANGE;
      return  NULL;
   }

   Comparable* found = NULL;
   int  increase = 0;

   cmp_t  result = root->Compare(item->Key());
   dir_t  dir = (result == MIN_CMP) ? LEFT : RIGHT;

   if (result != EQ_CMP) {
      // Insert into "dir" subtree 
      found = Insert(item, root->mySubtree[dir], change);
      if (found)  return  found;   // already here - dont insert
      increase = result * change;  // set balance factor increment
   } else  {   // key already in tree at this node
      increase = HEIGHT_NOCHANGE;
      return  root->myData;
   }

   root->myBal += increase;   

  // re-balance if needed 
   change =  (increase && root->myBal) ? (1 - ReBalance(root)) : HEIGHT_NOCHANGE;
   return  NULL;
}

Comparable* AvlNode::Delete(AvlKey key, AvlNode* &root, int &change, cmp_t cmp)
{
   if (root == NULL) {
      change = HEIGHT_NOCHANGE;
      return  NULL;
   }

   Comparable* found = NULL;
   int  decrease = 0;

   cmp_t  result = root->Compare(key, cmp);
   dir_t  dir = (result == MIN_CMP) ? LEFT : RIGHT;

   if (result != EQ_CMP) {
      found = Delete(key, root->mySubtree[dir], change, cmp);
      if (! found)  return  found;   // not found - can't delete
      decrease = result * change;    // set balance factor decrement
   } else  {   // Found key at this node
      found = root->myData;

      if ((root->mySubtree[LEFT] == NULL) &&
          (root->mySubtree[RIGHT] == NULL)) {
         // We have a leaf -- remove it
         delete  root;
         root = NULL;
         change = HEIGHT_CHANGE;// height changed from 1 to 0
         return  found;
      } else if ((root->mySubtree[LEFT] == NULL) ||
                 (root->mySubtree[RIGHT] == NULL)) {
         // We have one child -- only child becomes new root 
         AvlNode* toDelete = root;
         root = root->mySubtree[(root->mySubtree[RIGHT]) ? RIGHT : LEFT];
         change = HEIGHT_CHANGE;    // We just shortened the subtree
         // Null-out the subtree pointers so we dont recursively delete
         toDelete->mySubtree[LEFT] = toDelete->mySubtree[RIGHT] = NULL;
         delete  toDelete;
         return  found;
      } else {
         // We have two children -- find successor and replace our current
         // data item with that of the successor
         root->myData = Delete(key, root->mySubtree[RIGHT],
                               decrease, MIN_CMP);
      }
   }

   root->myBal -= decrease;//update balance   

   // Rebalance if necessary
   if (decrease) {
      if (root->myBal) {
         change = ReBalance(root);// rebalance and see if height changed
      } else {
         change = HEIGHT_CHANGE;// balanced because subtree decreased
      }
   } else {
      change = HEIGHT_NOCHANGE;
   }

   return  found;
}

int AvlNode::Height() const {
   int  leftHeight  = (mySubtree[LEFT])  ? mySubtree[LEFT]->Height()  : 0;
   int  rightHeight = (mySubtree[RIGHT]) ? mySubtree[RIGHT]->Height() : 0;
   return  (1 + MAX(leftHeight, rightHeight));
}

int AvlNode::Check() const {
   int  valid = 1;
   // First verify that subtrees are correct
   if (mySubtree[LEFT])   valid *= mySubtree[LEFT]->Check();
   if (mySubtree[RIGHT])  valid *= mySubtree[RIGHT]->Check();
   // Now get the height of each subtree
   int  leftHeight  = (mySubtree[LEFT])  ? mySubtree[LEFT]->Height()  : 0;
   int  rightHeight = (mySubtree[RIGHT]) ? mySubtree[RIGHT]->Height() : 0;
   // Verify that AVL tree property is satisfied
   int  diffHeight = rightHeight - leftHeight;
   if (LEFT_IMBALANCE(diffHeight) || RIGHT_IMBALANCE(diffHeight)) {
      valid = 0;
   }
   // Verify that balance-factor is correct
   if (diffHeight != myBal) {
      valid = 0;
   }
   // Verify that search-tree property is satisfied
   if ((mySubtree[LEFT]) && (mySubtree[LEFT]->Compare(Key()) == MIN_CMP)) {
      valid = 0;
   }
   if ((mySubtree[RIGHT]) && (mySubtree[RIGHT]->Compare(Key()) == MAX_CMP)) {
      valid = 0;
   }
   return  valid;
}

